#include "rmsnorm_kernel.h"
#include "math.h"
#include "algorithm"

template<typename T>
__device__ T WarpReduce(T val) {
    for (int i = 16; i >= 1; i /= 2) {
        val += __shfl_xor_sync(0xffffffff, val, i);
    }
    return val;
}

template<typename T>
__device__ T BlockReduce(T val) {
    int tid = threadIdx.x;
    int laneid = threadIdx.x % 32;
    int warpid = threadIdx.x / 32;
    int warp_num = (blockDim.x + 32 - 1) / 32;
    __shared__ T shared[64];
    T ret = WarpReduce<T>(val);
    if (laneid == 0) {
        shared[warpid] = ret;
    }
    __syncthreads();

    T sum = tid < warp_num ? shared[tid] : 0;
    sum = WarpReduce<T>(sum);
    return sum;
}

template<typename T>
__global__ void RMSNorm(T *decoder_out,
                        T *decoder_residual,
                        T *norm_weight,
                        float eps,
                        int hidden_units) 
{
    using Vec_t = typename Vec<T>::Type;
    int size = Vec<T>::size;
    int tid = threadIdx.x;
    int thread_sum = 0;
    Vec_t *cur_token_out = reinterpret_cast<Vec_t*>(decoder_out + blockIdx.x * hidden_units);
    Vec_t *cur_token_residual = reinterpret_cast<Vec_t*>(decoder_residual + blockIdx.x * hidden_units);
    for (int i = tid; i < hidden_units / size; i += blockDim.x) {
        cur_token_residual[i] = cur_token_out[i];
        Vec_t vec = cur_token_out[i];
        thread_sum += vec.x * vec.x;
        thread_sum += vec.y * vec.y;
        thread_sum += vec.z * vec.z;
        thread_sum += vec.w * vec.w;
    }
    float block_sum = BlockReduce<float>(thread_sum);
    __shared__ float inv_mean;
    if (threadIdx.x == 0) {
        inv_mean = rsqrtf(block_sum / hidden_units + eps);
    }
    __syncthreads();
    Vec_t *gamma_vec = reinterpret_cast<Vec_t*>(norm_weight);
    for (int i = tid; i < hidden_units / size; i += blockDim.x) {
        Vec_t vec = cur_token_out[i];
        Vec_t gamma = gamma_vec[i];
        cur_token_out[i].x = vec.x * inv_mean * gamma.x;
        cur_token_out[i].y = vec.y * inv_mean * gamma.y;
        cur_token_out[i].z = vec.z * inv_mean * gamma.z;
        cur_token_out[i].w = vec.w * inv_mean * gamma.w;
    }
}

template<>
__global__ void RMSNorm(half *decoder_out,
                        half *decoder_residual,
                        half *norm_weight,
                        float eps,
                        int hidden_units) 
{
    using Vec_t = typename Vec<half>::Type; 
    int size = Vec<half>::size;
    Vec_t *cur_token_out = reinterpret_cast<Vec_t*>(decoder_out + blockIdx.x * hidden_units);
    Vec_t *cur_token_residual = reinterpret_cast<Vec_t*>(decoder_residual + blockIdx.x * hidden_units);
    int tid = threadIdx.x;
    int thread_sum = 0;
    for (int i = tid; i < hidden_units / size; i += blockDim.x) {
        cur_token_residual[i] = cur_token_out[i];
        Vec_t vec = cur_token_out[i];
        thread_sum += __half2float(vec.x) * __half2float(vec.x);
        thread_sum += __half2float(vec.y) * __half2float(vec.y);
    }                          
    float block_sum = BlockReduce<float>(thread_sum);
    __shared__ float inv_mean;
    if (threadIdx.x == 0) {
        inv_mean = rsqrtf(block_sum / hidden_units + eps);
    }
    __syncthreads();
    Vec_t *gamma_vec = reinterpret_cast<Vec_t*>(norm_weight);
    for (int i = tid; i < hidden_units / size; i += blockDim.x) {
        Vec_t vec = cur_token_out[i];
        Vec_t gamma = gamma_vec[i];
        cur_token_out[i].x = gamma.x * __float2half(__half2float(vec.x) * inv_mean);
        cur_token_out[i].y = gamma.y * __float2half(__half2float(vec.y) * inv_mean);
    }
}

template<typename T>
void launchRMSNorm(TensorWrapper<T> *decoder_out,    // [num_tokens, hidden_units]
                   TensorWrapper<T> *decoder_residual,  // [num_tokens, hidden_units]
                   NormWeight<T> &norm_weight,   // [hidden_units, ]
                   float eps,
                   bool is_last) 
{
    size_t num_tokens = decoder_out->shape[0];
    int hidden_units = decoder_out->shape[1];
    dim3 grid(num_tokens);
    int num_threads = std::min(hidden_units/4, 1024);
    dim3 block(num_threads);
    RMSNorm<T><<<grid, block>>>(decoder_out->data,
                                decoder_residual->data,
                                norm_weight.gamma,
                                eps,
                                hidden_units);
}



template void launchRMSNorm(TensorWrapper<float> *decoder_out,    // [num_tokens, hidden_units]
                   TensorWrapper<float> *decoder_residual,  // [num_tokens, hidden_units]
                   NormWeight<float> &norm_weight,   // [hidden_units, ]
                   float eps,
                   bool is_last);

template void launchRMSNorm(TensorWrapper<half> *decoder_out,    // [num_tokens, hidden_units]
                   TensorWrapper<half> *decoder_residual,  // [num_tokens, hidden_units]
                   NormWeight<half> &norm_weight,   // [hidden_units, ]
                   float eps,
                   bool is_last);